import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

# ### 同步通信的阻塞特性
# - dist.send ：发送方（如进程0）调用 dist.send 后，会 阻塞等待 接收方（进程1）调用 dist.recv 准备好接收数据，否则发送方不会继续执行后续代码。
# - dist.recv ：接收方（如进程1）调用 dist.recv 后，会 阻塞等待 发送方（进程0）发送数据，否则接收方不会继续执行后续代码。
# 在你的代码中：

# - 进程0必须先执行 dist.send(tensor=1, dst=1) ，但发送操作会阻塞，直到进程1调用 dist.recv(src=0) 来接收数据。因此，进程1的 dist.recv 必须在进程0的 dist.send 之后才会完成，此时进程0的 tensor 已经被修改为1，进程1接收到的一定是1（而非初始的0）。
# - 同理，进程1修改 tensor 为2后调用 dist.send(dst=0) ，进程0的 dist.recv(src=1) 会阻塞等待，直到进程1的发送完成，因此进程0接收到的一定是2。
# 系统调度可能影响 打印语句的输出顺序 （例如进程1的 after recv 打印可能在进程0的 after send 打印之前或之后），但这是 日志输出的时间差 ，不影响通信的逻辑顺序。数据传输的核心操作（ send / recv ）始终按顺序完成，因此最终每个进程的 tensor 值一定是正确的。

# ### . 进程启动顺序与代码逻辑的关系
# 两个进程（ rank=0 和 rank=1 ）是并行启动的（通过 mp.Process 启动），系统调度可能让其中一个进程先进入 run 函数执行代码。但关键在于：

# - rank=0 的代码逻辑中， 主动发起了第一次通信 （调用 dist.send 发送数据）；
# - rank=1 的代码逻辑中， 被动等待接收通信 （调用 dist.recv 接收数据）。

# ### 2. 同步阻塞机制的强制顺序
# dist.send 和 dist.recv 是同步阻塞操作，它们的执行会互相等待：

# - 如果 rank=1 先启动并执行到 dist.recv ：此时 rank=1 的 recv 会阻塞，等待 rank=0 调用 send 发送数据（即使 rank=0 还未执行到 send ）。
# - 如果 rank=0 先启动并执行到 dist.send ：此时 rank=0 的 send 会阻塞，等待 rank=1 调用 recv 准备好接收数据（即使 rank=1 还未执行到 recv ）。

# ### dist.send 与 dist.recv 的同步阻塞机制
# - 发送方（ dist.send ） ： 当进程 A（如 rank=0 ）调用 dist.send(tensor, dst=B) 时，底层通信库（如 gloo ）会：
  
#   1. 向目标进程 B（ rank=1 ）发送一个“数据待接收”的请求（包含数据元信息，如张量大小、类型等）；
#   2. 阻塞等待进程 B 的响应（即进程 B 调用了 dist.recv(src=A) ）；
#   3. 一旦收到进程 B 的确认（表示 B 已准备好接收），发送方将实际数据（ tensor ）传输给 B，并解除阻塞，继续执行后续代码。
# - 接收方（ dist.recv ） ： 当进程 B（ rank=1 ）调用 dist.recv(tensor, src=A) 时，底层通信库会：
  
#   1. 监听来自进程 A 的“数据待接收”请求；
#   2. 收到请求后，向 A 发送确认（表示已准备好接收）；
#   3. 等待 A 传输实际数据，接收完成后解除阻塞，继续执行后续代码。

def run(rank_id, size):
    tensor = torch.zeros(1)
    if rank_id == 0:
        tensor += 1
        # Send the tensor to process 1
        dist.send(tensor=tensor, dst=1)
        print('after send, Rank ', rank_id, ' has data ', tensor[0])
        dist.recv(tensor=tensor, src=1)
        print('after recv, Rank ', rank_id, ' has data ', tensor[0])
    else:
        # Receive tensor from process 0
        dist.recv(tensor=tensor, src=0)
        print('after recv, Rank ', rank_id, ' has data ', tensor[0])
        tensor += 1
        dist.send(tensor=tensor, dst=0)
        print('after send, Rank ', rank_id, ' has data ', tensor[0])


def init_process(rank_id, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank_id, world_size=size)
    fn(rank_id, size)


if __name__ == "__main__":
    size = 2
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, run))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
